package comtgl.demo.utils;

import comtgl.demo.dto.RateLimitDto;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.stereotype.Component;

import java.util.concurrent.TimeUnit;

@Component
public class RateLimitUtils {

    @Autowired
    private RedisTemplate<String, Object> redisTemplate;

    /**
     * 检查是否可以执行
     * @param rateLimit 限流配置
     * @return 是否可以执行，true可以执行，false不可以
     */
    public boolean checkOK(RateLimitDto rateLimit) {
        // 生成限流key
        String secondKey = "rate_limit:second:" + rateLimit.getLimitKey();
        String minuteKey = "rate_limit:minute:" + rateLimit.getLimitKey();

        Object secondCount = redisTemplate.opsForValue().get(secondKey);
        if (secondCount == null) {
            Long currentSecondCount = redisTemplate.opsForValue().increment(secondKey, 1);
            if (currentSecondCount != null && currentSecondCount == 1) {
                redisTemplate.expire(secondKey, 1, TimeUnit.SECONDS);
            }
        }
        // 检查秒级限流
        secondCount = redisTemplate.opsForValue().get(secondKey);
        System.out.println("secondCount:" + secondCount);
        // 检查秒级限流
        Integer currentSecondCount = (Integer) secondCount;
        if (currentSecondCount != null && currentSecondCount > rateLimit.getSecondLimit()) {
            return false;
        }

        //分钟限制
        Object o1 = redisTemplate.opsForValue().get(minuteKey);
        if (o1 == null) {
            Long currentMinuteCount = redisTemplate.opsForValue().increment(minuteKey, 1);
            if (currentMinuteCount != null && currentMinuteCount == 1) {
                redisTemplate.expire(minuteKey, 1, TimeUnit.MINUTES);
            }
        }
        Object minuteCount = redisTemplate.opsForValue().get(minuteKey);
        System.out.println("minuteCount:" + minuteCount);
        // 检查分钟级限流
        Integer currentMinuteCount = (Integer) minuteCount;
        if (currentMinuteCount != null && currentMinuteCount > rateLimit.getMinuteLimit()) {
            return false;
        }
        return true;
    }

    public void incrementCount(String key) {
        String secondKey = "rate_limit:second:" + key;
        String minuteKey = "rate_limit:minute:" + key;
        long currentSecondCount = redisTemplate.opsForValue().increment(secondKey, 1);
        if (currentSecondCount == 1) {
            redisTemplate.expire(secondKey, 1, TimeUnit.SECONDS);
        }

        long currentMinuteCount = redisTemplate.opsForValue().increment(minuteKey, 1);
        if (currentMinuteCount == 1) {
            redisTemplate.expire(minuteKey, 1, TimeUnit.MINUTES);
        }
    }
}
